import cv2
from PIL import Image
import numpy as np
from tqdm import tqdm

from glob import glob
import os

from scipy.ndimage import gaussian_filter
from skimage.morphology import reconstruction
from skimage.segmentation import clear_border
from skimage.filters import threshold_otsu, threshold_sauvola
from skimage.measure import label, regionprops
from skimage.morphology import closing, square, dilation
from skimage.color import label2rgb
from skimage.transform import hough_line, hough_line_peaks, rotate

import numpy as np
from PIL import Image

# Pool Images
def poolImages(dir1, dir2, size):
    pathList = glob(os.path.join(dir1, '*'))
    for i in tqdm(range(len(pathList)), ncols=100):
        path1 = pathList[i]
        filename = os.path.split(path1)[-1]
        path2 = os.path.join(dir2, filename)
        image1 = Image.open(path1).convert('RGB')
        image2 = image1.resize(size, resample=Image.BOX)
        image2.save(path2)

def poolImages2(dir1, dir2, size):
    def poolImage2(image, size):
        src_size = np.array(image.size)
        dst_size = np.array(size)

        if image.mode == 'L':
            trans_array = np.zeros(size, dtype='uint8') + 255
        elif image.mode == 'RGB':
            trans_array = np.zeros((size[0], size[1], 3), dtype='uint8') + 255
        else:
            image = image.convert('RGB')
            trans_array = np.zeros((size[0], size[1], 3), dtype='uint8') + 255

        k = dst_size / src_size.max()
        scale_size = np.around((k * src_size)).astype('int')
        trans_vect = ((dst_size - scale_size) / 2).astype('int')

        scale_image = image.resize(tuple(scale_size), resample=Image.BOX)
        scale_array = np.array(scale_image, dtype='uint8')
        trans_array[trans_vect[1]: trans_vect[1]+scale_size[1],
                    trans_vect[0]: trans_vect[0]+scale_size[0]] =\
                    np.array(scale_image, dtype='uint8')

        return Image.fromarray(trans_array)

    pathList = glob(os.path.join(dir1, '*'))
    for i in tqdm(range(len(pathList)), ncols=100):
        path1 = pathList[i]
        filename = os.path.split(path1)[-1]
        path2 = os.path.join(dir2, filename)
        image = poolImage2(Image.open(path1), size)
        image.save(path2)

# Calc boxs
def calcLabelImage(image):
    cleared = clear_border(image)
    label_image = label(cleared)
    return label_image

def calcAreaBoxs(predict_y, limit=10):
    threshold = predict_y.max() - 0.9 * (predict_y.max() - predict_y.min())
    image = (predict_y > threshold).astype('uint8')

    label_image = calcLabelImage(image)

    boxs = []
    for region in regionprops(label_image):
        if region.area >= limit:
            minr, minc, maxr, maxc = region.bbox
            boxs.append([[minr, minc], [maxr, maxc]])
    return boxs

def calcSourceBoxs(image, boxs, a):
    if len(boxs) == 0: return np.array([])

    w, h = image.size
    size = np.array([h, w])
    k = np.max(size) / a

    boxs_arr = np.array(boxs)
    boxs_arr = boxs_arr * k + (size - np.max(size)) / 2
    # return np.maximum(0, np.round(boxs_arr)).astype('int')
    boxs_arr = np.maximum(0, np.round(boxs_arr)).astype('int')
    return np.array([box_arr for box_arr in boxs_arr if (box_arr[1, 0] <= h and box_arr[1, 1] <= w)])

def filterImage(img_arr):
    image = 1 - img_arr.astype('float') / 255
    image = gaussian_filter(image, 1)
    seed = np.copy(image)
    # Debug code.
    # print('Debug: {}'.format(seed.shape))
    seed[1:-1, 1:-1] = image.min()
    mask = image
    dilated = reconstruction(seed, mask, method='dilation')
    return image - dilated

def labelImage(img_arr, l=100):
    # apply threshold
    if img_arr.max() - img_arr.min() < 0.1:
        bw = img_arr
    else:
        thresh_o = threshold_otsu(img_arr)
        window_size = 25
        thresh_s = threshold_sauvola(img_arr, window_size=window_size)
        bw = dilation(((img_arr > thresh_o) & (img_arr > thresh_s)), square(2))
        bw = closing(bw, square(3))

    # remove artifacts connected to image border
    cleared = clear_border(bw)

    # label image regions
    label_image = label(cleared)
    
    bboxs = []
    for region in regionprops(label_image):
        if region.area > l:
            bboxs.append(region.bbox)

    return np.array(bboxs).reshape(-1, 2, 2)

def filterBBoxs(bboxs):
    l = np.sort((bboxs[:, 1, :] - bboxs[:, 0, :]).flatten())[len(bboxs)] * 1.5
    pts = []
    for pt0, pt1 in bboxs:
        nk0 = (np.sum((pt0 - bboxs[:, 0, :])**2, axis=1)**0.5 < l).sum()
        nk1 = (np.sum((pt1 - bboxs[:, 1, :])**2, axis=1)**0.5 < l).sum()
        if (nk0 + nk1) > 2: pts.append([pt0, pt1])
    return np.array(pts).reshape(-1, 2, 2)

def sliceH(H):
    slice_pts = []
    flag = True
    for i, h in enumerate(H):
        if (h > 0) == flag:
            slice_pts.append(i)
            flag = not flag
    else:
        if len(slice_pts) % 2 != 0:
            slice_pts.append(i)

    return slice_pts

def packetBBoxs(bboxs, shape):
    pts = bboxs.reshape(-1, 4)
    packet_bboxs = []
    if shape[0] <= shape[1]:
        h = np.zeros((shape[0]))
        for minr, _, maxr, _ in pts:
            h[minr: maxr] = 1
        sliceh = sliceH(h)

        for gminr, gmaxr in np.array(sliceh).reshape(-1, 2):
            M = (pts[:, 0] >= gminr) * (pts[:, 2] <= gmaxr)
            gminc, gmaxc = shape[1], 0
            for _, minc, _, maxc in pts[M]:
                gminc = minc if minc < gminc else gminc
                gmaxc = maxc if maxc > gmaxc else gmaxc
            packet_bboxs.append([gminr, gminc, gmaxr, gmaxc])
    else:
        h = np.zeros((shape[1]))
        for _, minc, _, maxc in pts:
            h[minc: maxc] = 1
        sliceh = sliceH(h)

        for gminc, gmaxc in np.array(sliceh).reshape(-1, 2):
            M = (pts[:, 1] >= gminc) * (pts[:, 3] <= gmaxc)
            gminr, gmaxr = shape[0], 0
            for minr, _, maxr, _ in pts[M]:
                gminr = minr if minr < gminr else gminr
                gmaxr = maxr if maxr > gmaxr else gmaxr
            packet_bboxs.append([gminr, gminc, gmaxr, gmaxc])
    return np.array(packet_bboxs).reshape(-1, 2, 2)

def calcAccurateBoxs(image_arr, boxs_arr, l=0.3):
    boxs_arr_int = np.maximum(0, boxs_arr).astype('int')

    accurate_boxs = []
    for box_arr in boxs_arr_int:
        if np.min(box_arr[1, :] - box_arr[0, :]) < 1: continue
        sub_area = image_arr[box_arr[0, 0]:box_arr[1, 0], box_arr[0, 1]:box_arr[1, 1]]
        # Test code.
        if sub_area.shape[0] <=1 or sub_area.shape[1] <= 1: continue
        img = filterImage(sub_area)
        bboxs = labelImage(img)
        if len(bboxs) < 1: continue
        if len(bboxs) > 2: filterBBoxs(bboxs)
        bboxs = packetBBoxs(bboxs, shape=img.shape)
        bboxs[:, :, 0] += box_arr[0, 0]
        bboxs[:, :, 1] += box_arr[0, 1]
        accurate_boxs.extend(bboxs.tolist())

    if len(accurate_boxs) == 0: return boxs_arr
    return np.array(accurate_boxs)

# TextCleaner
def active(X, W):
    M = W * (1 - X.astype('float32') / 255)
    C = 1 - M * (M > 0.25)
    A = np.tan(np.pi / 3 * C - np.pi / 12)
    return (np.maximum(A, 0) * 255).astype('uint8')

def textCleaner(input_image, output_path):
    if isinstance(input_image, Image.Image):
        img = np.array(input_image.convert('L'), dtype='uint8')
    else:
        img = cv2.imread(input_image, 0)

    h, w = img.shape[:2]
    k = 600 / max(h, w)
    mini = cv2.resize(img, (round(w / 4), round(h / 4)))

    kernel = np.ones((3, 3))

    edges = cv2.Canny(mini, 100, 200)

    morphology = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel, iterations=1)
    median = cv2.medianBlur(morphology, 3)

    morphology = cv2.morphologyEx(median, cv2.MORPH_OPEN, kernel, iterations=1)
    morphology = cv2.morphologyEx(morphology, cv2.MORPH_DILATE, kernel, iterations=2)

    kernel = np.ones((5, 5)) / 9
    filted = cv2.filter2D(morphology, -1, kernel)

    mask = filted.astype('float32') / filted.max()
    mask = cv2.resize(mask, (w, h))
    clean = active(img, mask)
    cv2.imwrite(output_path, clean)

# TextDeskew
def calculate_deviation(angle):
    angle_in_degrees = np.abs(angle)
    deviation = np.abs(np.pi / 4 - angle_in_degrees)

    return deviation

def compare_sum(value):
    if value >= 44 and value <= 46:
        return True
    else:
        return False

def get_max_freq_elem(arr):
    max_arr = []
    freqs = {}
    for i in arr:
        if i in freqs:
            freqs[i] += 1
        else:
            freqs[i] = 1

    sorted_keys = sorted(freqs, key=freqs.get, reverse=True)
    max_freq = freqs[sorted_keys[0]]

    for k in sorted_keys:
        if freqs[k] == max_freq:
            max_arr.append(k)

    return max_arr

def textDeskew(input_image, output_path, model, limit=1, no_bias=False, format_size=(128, 128)):
    if isinstance(input_image, Image.Image):
        if input_image.mode == 'RGB':
            img = np.array(Image.merge('RGB', input_image.split()[::-1]), dtype='uint8')
        else:
            img = np.array(input_image, dtype='uint8')
    else:
        img = cv2.imread(input_image)

    if len(img.shape) == 3:
        # Test code.
        # gray = img.mean(axis=2, dtype='uint8') # cv2.imread(input_image, 0)
        gray = img[:,:, :3].mean(axis=2, dtype='uint8') # cv2.imread(input_image, 0)
    else:
        gray = img

    bias_angle = 0
    if not no_bias:
        size = np.array(gray.shape[::-1])
        sample_k = 2048 / size.max()
        sample_size = np.around(sample_k * size).astype('int')
        sample = cv2.resize(gray, tuple(sample_size), interpolation=cv2.INTER_AREA)

        edges = cv2.Canny(sample, 50, 150, apertureSize=3)

        h, a, d = hough_line(edges)
        _, ap, _ = hough_line_peaks(h, a, d, num_peaks=20)

        if len(ap) != 0:
            absolute_deviations = [calculate_deviation(k) for k in ap]
            average_deviation = np.mean(np.rad2deg(absolute_deviations))
            ap_deg = [np.rad2deg(x) for x in ap]

            bin_0_45 = []
            bin_45_90 = []
            bin_0_45n = []
            bin_45_90n = []
            
            for ang in ap_deg:
                deviation_sum = int(90 - ang + average_deviation)
                if compare_sum(deviation_sum):
                    bin_45_90.append(ang)
                    continue

                deviation_sum = int(ang + average_deviation)
                if compare_sum(deviation_sum):
                    bin_0_45.append(ang)
                    continue

                deviation_sum = int(-ang + average_deviation)
                if compare_sum(deviation_sum):
                    bin_0_45n.append(ang)
                    continue

                deviation_sum = int(90 + ang + average_deviation)
                if compare_sum(deviation_sum):
                    bin_45_90n.append(ang)

            angles = [bin_0_45, bin_45_90, bin_0_45n, bin_45_90n]
            lmax = 0

            for j in range(len(angles)):
                l = len(angles[j])
                if l > lmax:
                    lmax = l
                    maxi = j

            if lmax:
                ans_arr = get_max_freq_elem(angles[maxi])
                ans_res = np.mean(ans_arr)

            else:
                ans_arr = get_max_freq_elem(ap_deg)
                ans_res = np.mean(ans_arr)

            angle = ans_res

            if angle >= 0 and angle <= 90:
                bias_angle = angle - 90
            if angle >= -45 and angle < 0:
                bias_angle = angle - 90
            if angle >= -90 and angle < -45:
                bias_angle = 90 + angle

        rotated1 = rotate(img, bias_angle, resize=True, cval=1)
    else:
        rotated1 = img.astype('float') / 255

    mini = np.ones(format_size)
    rotated1_size = np.array([rotated1.shape[1], rotated1.shape[0]])
    pool_k = 128 / rotated1_size.max()
    pool_size = np.around(pool_k * rotated1_size).astype('int')
    pool = cv2.resize(rotated1, tuple(pool_size), interpolation=cv2.INTER_AREA)
    if len(pool.shape) > 2: pool = pool.mean(axis=2)
    mini_limits = (pool_size - pool_size.min()) // 2
    mini[mini_limits[0]: mini_limits[0]+pool_size[1], mini_limits[1]: mini_limits[1]+pool_size[0]] = pool
    input_X = 1 - mini.reshape(-1, format_size[0], format_size[1], 1).astype('float32')
    predict_probs = model.predict(input_X)
    direction_num = (predict_probs == predict_probs.max()).dot(np.arange(4))

    rotated2 = (np.rot90(rotated1, k=-direction_num) * 255).astype('uint8')

    cv2.imwrite(output_path, rotated2)

    print('{}: Deskew Angle: {:.2f}'.format(input_image, bias_angle + -90 * int(direction_num)))
    return rotated2
